﻿using System;
using Telerik.JustDecompiler.Cil;
using Telerik.JustDecompiler.Decompiler;
using Telerik.JustDecompiler.Decompiler.StateMachines;

namespace Telerik.JustDecompiler.Steps
{
    /// <summary>
    /// Removes the yeild state machine from the CFG of the method.
    /// </summary>
    /// <remarks>
    /// A very good article that explains the yield state machine: "C# in Depth  Iterator block implementation details.htm".
    /// All of the explanations below presume that the reader has read the mentioned article.
    /// 
    /// The general idea of this step is to remove the state machine from the CFG and to rearrange the successor-predecessor relations between
    /// the instruction blocks of the CFG, so that the control flow of the resulting CFG to be more similar to the original code.
    /// Also this step marks up the instruction blocks that should contain 'yield return' and 'yield break' as expressions.
    /// Finally we get the exception data, from the Dispose method, that is going to be used by the LogicalFlowBuilder to recreate the removed try/finally
    /// constructs.
    /// 
    /// UPDATE:
    /// The newer versions of the C# compiler generate different state machine. Try/finally constructs are no longer removed. Instead a state controller is
    /// inserted at the begining of each try/finally block. Execution of finally handlers is controlled by a new compiler generated variable (see
    /// StateMachineDoFinallyCheckRemover class for detailed explanation). A new field is also generated by the compiler. This field is set by the Dispose method
    /// and is used when disposing the yield iterator (see DisposingStateControllerRemover for details).
    /// 
    /// Since there are two versions of the yield state machine, here is how the state machine is removed (old state machine - V1, new - V2):
    /// 1) Analyze the dispose method - The analysis gives us the version of the state machine and the exception handler data for the V1 state machine, or
    /// the disposing field for V2 state machine
    /// 2) Determine the state controller remover based on the version of the state machine - if it's V2, then run an additional step to mark for removal
    /// the do-finally-check blocks
    /// 3) Mark the blocks of the state controller for removal and find the entry of each state
    /// 4) Attach the end of each state to the begining of the next state and determine the blocks that contain yield return or yield break
    /// 5) Remove the blocks that were marked for removal
    /// </remarks>
    internal class RemoveYieldStateMachineStep : BaseStateMachineRemoverStep
    {
        protected override bool ProcessCFG()
        {
            StateMachineDisposeAnalyzer disposeAnalyzer = new StateMachineDisposeAnalyzer(moveNextMethodContext.Method);
            StateControllerRemover controllerRemover;
            YieldStateMachineVersion machineVersion = disposeAnalyzer.ProcessDisposeMethod();
            if (machineVersion == YieldStateMachineVersion.V1)
            {
                controllerRemover = new StateControllerRemover(moveNextMethodContext);
            }
            else if (machineVersion == YieldStateMachineVersion.V2)
            {
                controllerRemover = new DisposingStateControllerRemover(moveNextMethodContext, disposeAnalyzer.StateField, disposeAnalyzer.DisposingField);
                StateMachineDoFinallyCheckRemover doFinallyCheckRemover = new StateMachineDoFinallyCheckRemover(moveNextMethodContext);
                if (!doFinallyCheckRemover.MarkFinallyConditionsForRemoval())
                {
                    return false;
                }
                toBeRemoved.UnionWith(doFinallyCheckRemover.BlocksMarkedForRemoval);
            }
            else
            {
                return false;
            }

            if (!controllerRemover.RemoveStateMachineController())
            {
                return false;
            }

            toBeRemoved.UnionWith(controllerRemover.BlocksMarkedForRemoval);

            SwitchData switchData = controllerRemover.SwitchData;
            YieldStateMachineControlFlowRebuilder controlFlowRebuilder =
                new YieldStateMachineControlFlowRebuilder(moveNextMethodContext, switchData, controllerRemover.StateField);

            if (!controlFlowRebuilder.ProcessEndBlocks())
            {
                return false;
            }

            toBeRemoved.UnionWith(controlFlowRebuilder.BlocksMarkedForRemoval);
            StateMachineCFGCleaner cfgCleaner = new StateMachineCFGCleaner(theCFG, switchData, switchData.OrderedCasesArray[0]);

            if (!cfgCleaner.CleanUpTheCFG(toBeRemoved))
            {
                return false;
            }

            YieldFieldsInformation fieldsInfo = new YieldFieldsInformation(controllerRemover.StateField,
                    controlFlowRebuilder.CurrentItemField, controlFlowRebuilder.ReturnFlagVariable);

            this.moveNextMethodContext.YieldData =
                new YieldData(machineVersion, controlFlowRebuilder.YieldReturnBlocks,
                    controlFlowRebuilder.YieldBreakBlocks, fieldsInfo, disposeAnalyzer.YieldsExceptionData);
            return true;
        }
    }
}
